package com.coderworm.tcp2;
import java.io.*;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.nio.ByteOrder;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.zip.CRC32;

/**
 * 带用户名密码验证的RTCM TCP服务端
 * 安全特性：身份验证（SHA-256加盐）、IP白名单、连接频率限制、数据合法性校验
 */
public class RtcmTcpServer2 {
    // 1. 基础配置
    private static final int TCP_PORT = 8887; // 监听端口
    private static final String RTCM_SAVE_DIR = "D:/rtcm_files/"; // 文件保存路径
    private static final String SALT = "RtcmSecure2025_#"; // 密码加盐值（客户端需相同）

    // 2. 安全配置（IP白名单、连接频率限制）
    private static final Set<String> ALLOWED_IPS = new HashSet<>(Arrays.asList(
            "127.0.0.1", "192.168.1.100", "192.168.1.101" // 授权设备IP
    ));
    private static final int CONNECT_LIMIT = 5; // 单IP最大连接次数（每分钟）
    private static final Map<String, AtomicInteger> IP_CONNECT_COUNT = new ConcurrentHashMap<>();
    private static final Timer CONNECT_TIMER = new Timer();
    private static final   ExecutorService threadPool = Executors.newFixedThreadPool(2);

    // 3. 授权用户信息（存储用户名+加盐后的密码SHA-256值，避免明文存储）
    private static final Map<String, String> AUTH_USERS = new HashMap<>();
    static {
        // 示例用户：user=rtcm_device_001, pwd=SecurePwd@2025 → 计算加盐后的SHA-256
        AUTH_USERS.put("rtcm_device_001",
                "b39c8fa70ebd421d39ed55d49c6d6d075ecfa075f0d00f2b2f687142d73aeff5");
        AUTH_USERS.put("rtcm_device_002",
                "a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b2");
    }

    public static void main(String[] args) {
        // 初始化连接计数器定时清理任务（每分钟重置一次）
        CONNECT_TIMER.scheduleAtFixedRate(new TimerTask() {
            @Override
            public void run() {
                IP_CONNECT_COUNT.clear();
                System.out.println("连接计数器已重置（每分钟清理）");
            }
        }, 0, 60 * 1000);

        // 启动服务端
        try (ServerSocket serverSocket = new ServerSocket(TCP_PORT);) {

            System.out.println("带身份验证的RTCM TCP服务端已启动，监听端口：" + TCP_PORT);
            System.out.println("授权IP列表：" + ALLOWED_IPS);

            while (true) {
                Socket clientSocket = serverSocket.accept();
                String clientIp = clientSocket.getInetAddress().getHostAddress();
                System.out.println("\n收到连接请求：IP=" + clientIp);

                // 提交任务到线程池处理
                threadPool.submit(new RtcmFileHandler(clientSocket, clientIp));
            }
        } catch (IOException e) {
            System.err.println("服务端启动异常：" + e.getMessage());
            e.printStackTrace();
        } finally {
            CONNECT_TIMER.cancel();
        }
    }

    /**
     * 文件处理任务类（包含身份验证、文件接收逻辑）
     */
    static class RtcmFileHandler implements Runnable {
        private final Socket clientSocket;
        private final String clientIp;

        public RtcmFileHandler(Socket clientSocket, String clientIp) {
            this.clientSocket = clientSocket;
            this.clientIp = clientIp;
        }

        @Override
        public void run() {
            System.out.println("[Handler-run] 开始执行run方法，处理客户端：IP=" + clientIp);
            try (InputStream in = clientSocket.getInputStream();
                 BufferedInputStream bis = new BufferedInputStream(in);
                 OutputStream out = clientSocket.getOutputStream();
                 BufferedOutputStream bos = new BufferedOutputStream(out)) {
                System.out.println("获取客户端输入输出流。。。。" );
                // -------------------------- 安全校验1：IP白名单 --------------------------
                if (!ALLOWED_IPS.contains(clientIp)) {
                    sendResponse(bos, "ERROR: 未授权IP，拒绝连接");
                    System.err.println("非法IP连接：" + clientIp);
                    return;
                }

                // -------------------------- 安全校验2：连接频率限制 --------------------------
                AtomicInteger connectCount = IP_CONNECT_COUNT.computeIfAbsent(clientIp, k -> new AtomicInteger(0));
                if (connectCount.incrementAndGet() > CONNECT_LIMIT) {
                    sendResponse(bos, "ERROR: 连接频率过高（每分钟最多" + CONNECT_LIMIT + "次），暂时拒绝");
                    System.err.println("IP连接超限：" + clientIp + "（次数：" + connectCount.get() + "）");
                    return;
                }

                // -------------------------- 核心：身份验证（用户名+密码密文） --------------------------
                if (!authenticate(bis, bos)) {
                    return; // 验证失败，直接断开
                }

                // -------------------------- 原有逻辑：文件接收与保存 --------------------------
                // 1. 解析文件元数据（文件名长度、文件名、文件大小、CRC）
                String fileName = parseFileName(bis, bos);
                if (fileName == null) return;

                long fileTotalSize = parseFileSize(bis, bos);
                if (fileTotalSize <= 0 || fileTotalSize > 1024 * 1024 * 1024) { // 限制1GB以内
                    sendResponse(bos, "ERROR: 非法文件大小（需1~1024MB）");
                    return;
                }

                long expectedCrc = parseCrc(bis, bos);
                if (expectedCrc < 0) {
                    sendResponse(bos, "ERROR: CRC校验码读取失败");
                    return;
                }

                // 2. 校验RTCM文件合法性（首字节必须为0xD3）
                if (!validateRtcmFileHeader(bis)) {
                    sendResponse(bos, "ERROR: 非法RTCM文件（缺少0xD3帧头）");
                    return;
                }

                // 3. 接收并保存文件
                File rtcmFile = new File(RTCM_SAVE_DIR + fileName);
                if (saveRtcmFile(bis, rtcmFile, fileTotalSize, expectedCrc)) {
                    sendResponse(bos, "SUCCESS: 文件接收完成，保存路径：" + rtcmFile.getAbsolutePath());
                    System.out.println("文件接收成功：IP=" + clientIp + "，文件=" + fileName);
                } else {
                    sendResponse(bos, "ERROR: 文件接收失败（CRC校验不通过或写入错误）");
                    if (rtcmFile.exists()) rtcmFile.delete(); // 删除损坏文件
                }

            } catch (Exception e) {
                System.err.println("处理异常（IP=" + clientIp + "）：" + e.getMessage());
                e.printStackTrace();
            } finally {
                // 关闭连接
                try {
                    if (!clientSocket.isClosed()) clientSocket.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
                System.out.println("连接已关闭：IP=" + clientIp);
            }
        }

        /**
         * 身份验证：接收客户端用户名+密码密文，校验合法性
         */
        private boolean authenticate(BufferedInputStream bis, BufferedOutputStream bos) throws IOException {
            // 1. 接收用户名（长度2字节+内容）
            String userName = receiveString(bis, bos, "用户名");
            if (userName == null) return false;

            // 2. 接收密码密文（长度2字节+内容，SHA-256字符串）
            String pwdHash = receiveString(bis, bos, "密码密文");
            if (pwdHash == null) return false;

            // 3. 校验用户名和密码密文
            String authPwdHash = AUTH_USERS.get(userName);
            if (authPwdHash == null || !authPwdHash.equals(pwdHash)) {
                sendResponse(bos, "ERROR: 用户名或密码错误，拒绝连接");
                System.err.println("身份验证失败：IP=" + clientIp + "，用户名=" + userName);
                return false;
            }

            // 验证通过
            sendResponse(bos, "AUTH_SUCCESS: 身份验证通过，允许文件传输");
            System.out.println("身份验证成功：IP=" + clientIp + "，用户名=" + userName);
            return true;
        }

        /**
         * 接收字符串（长度2字节+内容，UTF-8编码）
         */
        private String receiveString(BufferedInputStream bis, BufferedOutputStream bos, String dataName) throws IOException {
            // 读取长度（2字节，大端序）
            byte[] lenBytes = new byte[2];
            int readLen = bis.read(lenBytes);
            if (readLen != 2) {
                sendResponse(bos, "ERROR: " + dataName + "长度读取失败");
                return null;
            }
            int strLen = ByteBuffer.wrap(lenBytes).order(ByteOrder.BIG_ENDIAN).getShort() & 0xFFFF;

            // 读取内容
            byte[] strBytes = new byte[strLen];
            readLen = bis.read(strBytes);
            if (readLen != strLen) {
                sendResponse(bos, "ERROR: " + dataName + "内容读取失败");
                return null;
            }

            return new String(strBytes, StandardCharsets.UTF_8);
        }

        /**
         * 解析文件名
         */
        private String parseFileName(BufferedInputStream bis, BufferedOutputStream bos) throws IOException {
            return receiveString(bis, bos, "文件名");
        }

        /**
         * 解析文件大小（8字节，大端序）
         */
        private long parseFileSize(BufferedInputStream bis, BufferedOutputStream bos) throws IOException {
            byte[] sizeBytes = new byte[8];
            int readLen = bis.read(sizeBytes);
            if (readLen != 8) {
                sendResponse(bos, "ERROR: 文件大小读取失败");
                return -1;
            }
            return ByteBuffer.wrap(sizeBytes).order(ByteOrder.BIG_ENDIAN).getLong();
        }

        /**
         * 解析CRC32校验码（4字节，大端序）
         */
        private long parseCrc(BufferedInputStream bis, BufferedOutputStream bos) throws IOException {
            byte[] crcBytes = new byte[4];
            int readLen = bis.read(crcBytes);
            if (readLen != 4) {
                sendResponse(bos, "ERROR: CRC校验码读取失败");
                return -1;
            }
            return ByteBuffer.wrap(crcBytes).order(ByteOrder.BIG_ENDIAN).getInt() & 0xFFFFFFFFL;
        }

        /**
         * 校验RTCM文件首字节（必须为0xD3）
         */
        private boolean validateRtcmFileHeader(BufferedInputStream bis) throws IOException {
            // 关键1：标记当前流位置，允许后续读取1字节后重置（readlimit=1）
            // 注意：BufferedInputStream默认缓冲区大小≥1，此处标记有效
            bis.mark(1);
            byte[] header = new byte[1];
            int readLen = bis.read(header);
            if (readLen != 1 || header[0] != (byte) 0xD3) {
                return false;
            }
            // 把读取的帧头放回流中（后续保存文件需要）
            bis.reset();
            System.out.println("[帧头校验] 合法RTCM帧头：0x" + Integer.toHexString(header[0] & 0xFF).toUpperCase());
            return true;
        }

        /**
         * 保存RTCM文件并校验CRC
         */
        private boolean saveRtcmFile(BufferedInputStream bis, File file, long totalSize, long expectedCrc) throws IOException {
            // 创建保存目录
            File parentDir = file.getParentFile();
            if (!parentDir.exists() && !parentDir.mkdirs()) {
                System.err.println("目录创建失败：" + parentDir.getAbsolutePath());
                return false;
            }

            // 接收文件并计算CRC
            CRC32 crc32 = new CRC32();
            byte[] buffer = new byte[4096];
            long receivedSize = 0;

            try (FileOutputStream fos = new FileOutputStream(file);
                 BufferedOutputStream fileBos = new BufferedOutputStream(fos)) {

                while (receivedSize < totalSize) {
                    int maxRead = (int) Math.min(buffer.length, totalSize - receivedSize);
                    int readBytes = bis.read(buffer, 0, maxRead);

                    if (readBytes == -1) {
                        System.err.println("文件接收中断：已接收" + receivedSize + "/" + totalSize + "字节");
                        return false;
                    }

                    fileBos.write(buffer, 0, readBytes);
                    crc32.update(buffer, 0, readBytes);
                    receivedSize += readBytes;

                    // 打印进度（每10%更新一次）
                    int progress = (int) ((receivedSize * 100.0) / totalSize);
                    if (progress % 10 == 0 && progress != (int) (((receivedSize - readBytes) * 100.0) / totalSize)) {
                        System.out.println("接收进度（IP=" + clientIp + "）：" + progress + "%");
                    }
                }
                fileBos.flush();
            }

            // 校验CRC
            long actualCrc = crc32.getValue();
            return actualCrc == expectedCrc;
        }

        /**
         * 发送响应给客户端（UTF-8编码，末尾加换行符）
         */
        private void sendResponse(BufferedOutputStream bos, String response) throws IOException {
            byte[] responseBytes = (response + "\n").getBytes(StandardCharsets.UTF_8);
            bos.write(responseBytes);
            bos.flush();
            System.out.println("发送响应（IP=" + clientIp + "）：" + response);
        }
    }

    /**
     * 辅助方法：计算字符串的SHA-256哈希值（加盐）
     * 客户端需使用相同方法计算密码密文
     */
    public static String calculateSha256(String content, String salt) {
        try {
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            byte[] input = (content + salt).getBytes(StandardCharsets.UTF_8);
            byte[] hashBytes = digest.digest(input);

            // 转为16进制字符串
            StringBuilder sb = new StringBuilder();
            for (byte b : hashBytes) {
                sb.append(String.format("%02x", b));
            }
            return sb.toString();
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("SHA-256算法不存在", e);
        }
    }
}